import time
import numpy as np
import logging
import itertools

def ddmCodeSel(tx, slot):
    elements = range(1, slot)
    comb = tx - 1
    comb_list = np.array([(0,)+i for i in itertools.combinations(elements, comb)])
    k = np.array([(comb_list+i)%slot for i in range(1, slot)])
    maxHit = []
    for i in range(len(comb_list)):
        union = [len(np.intersect1d(j, comb_list[i,:])) for j in k[:,i,:]]
        maxHit.append(np.max(union))
    minHitVal = np.min(maxHit)
    selComb = comb_list[np.argwhere(maxHit==minHitVal).flatten()]
    code = np.sum(0x1<<selComb, axis=-1)
    code = ['0b' + bin(i)[2:].zfill(slot) for i in code]
    return minHitVal, selComb, code
    
def stamp2str(x): return time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(x))


def str2stamp(x): return time.mktime(time.strptime(x, '%Y-%m-%d %H:%M:%S'))


def dB(voltage: np.ndarray): return 20*np.log10(voltage+0.00001)


def log2dB(x: np.ndarray): return x / 85


def dB2log(x: np.ndarray): return x * 85


def deCompress(M_in: bytearray, data_packing='CP4D', endians='big'):
    """de compress nxp`s cp8 or cp4 data

    Args:
        M_in (bytearray): input bytes of compress data
        data_packing (str, optional): CP4D or CP8. Defaults to 'CP4D'.
        endians (str, optional): endians. Defaults to 'big'.

    Returns:
        np.ndarray: decompress data
    """
    M_in = np.frombuffer(M_in, dtype='>Q' if endians == 'big' else 'Q')
    if len(M_in) % 8 != 0:
        logging.warning(f"Input data length{len(M_in)}illegality!")
        return None
    common = M_in >> 60 if endians == 'big' else (M_in & 0xF)
    common = common.astype('H')
    if data_packing == 'CP8':
        indices = np.arange(16)
        effBits = np.r_[3 * np.ones(2), 4 * np.ones(12), 3 * np.ones(2)].astype('H')
        mask = 2**effBits - 1
        mantissaPtr = 60 if endians == 'big' else 4
    else:
        indices = np.arange(8)
        effBits = 7 * np.ones(8, dtype='H')
        mask = 2**effBits - 1
        mantissaPtr = 56 if endians == 'big' else 8
    o = np.empty([len(effBits), len(M_in)], dtype='I')
    for e, m, i in zip(effBits, mask, indices):
        if endians == 'big':
            mantissaPtr -= e
            o[i, :] = ((M_in >> mantissaPtr) & m) << (32 - e)
        else:
            o[i, :] = ((M_in >> mantissaPtr) & m) << (32 - e)
            mantissaPtr += e
            
    o.dtype = np.int32
    o = o >> (8 + common)
    o = o.T.reshape(-1).astype(np.float32)
    o.dtype = np.complex64
    return o

def cmplx2int16(inArray: np.ndarray):
    array = inArray.ravel().astype(np.complex128)
    array.dtype = np.float64
    array = array / np.max(abs(array)) * 0x7FFF
    array = np.round(array, 0).astype('h')
    return array

def SptHist(mag: np.ndarray, inputType='real'):
    bins = 256*np.arange(1, 64)  # log2 coder
    ref = bins[1:] - 128
    if inputType == 'real':
        mag = np.round(20*np.log10(mag)*85).astype('H')
        ref = 10**(ref / 85 / 20)
    elif inputType == 'dB':
        mag = (mag * 85).astype('H')
        ref = ref / 85
    else:
        mag = mag.astype('H')
    ret = []
    for i in mag:
        tmp = np.histogram(i, bins)
        ret.append(ref[np.argmax(tmp[0])])
    ret = np.array(ret)
    return ret

def hybridCfar(cuberaw, cube3d, ddmOrder, leakOrder, cfarThld):
    rBin = cube3d.shape[0]
    dBin = cube3d.shape[1]
    aBin = cube3d.shape[2]
    ddmBank = len(ddmOrder) + len(leakOrder)
    ddmDist = dBin // ddmBank
    nccMag = np.mean(abs(cuberaw), axis=-1)
    mag = abs(cube3d)
    magZip = np.mean(mag.reshape([rBin, ddmBank, ddmDist, aBin]), axis=1)
    thld = np.sort(magZip, axis=1)
    thld = thld[:,int(0.5*ddmDist),:]
    peak = np.zeros(magZip.shape, dtype='b')
    for i in range(aBin):
        peak[:,:,i] = np.where(magZip[:,:,i].T > thld[:,i]*(10**(cfarThld/20)), 1, 0).T
    peak = np.sum(peak, axis = -1)
    idx = np.argwhere(peak != 0)
    aziBins = np.argmax(magZip[idx[:,0], idx[:,1], :],axis=-1)
    ddmVal = np.array([nccMag[idx[:,0], idx[:,1]+i*ddmDist] for i in range(ddmBank)]).T
    ddmCost = [np.sum(ddmVal[:,(ddmOrder + i)%ddmBank], axis=-1) for i in range(ddmBank)]
    ddmCost = np.array(ddmCost).T
    ddmRet = np.argmax(ddmCost, axis=-1)
    pwr = np.max(ddmCost ,axis=-1)
    conf1 = pwr / np.sum(ddmVal, axis=-1)
    conf2 = np.zeros(len(conf1))
    snr = 20*np.log10(pwr / thld[idx[:,0], aziBins])
    for i in range(len(ddmRet)):
        _peak = ddmVal[i, (ddmOrder + ddmRet[i])%ddmBank].min()
        _leak = ddmVal[i, (leakOrder + ddmRet[i])%ddmBank].max()
        conf2[i] = _peak / _leak
    # vaild = np.argwhere((conf1 > 0.8) & (conf2 > 1.3)).flatten()
    logic = (conf1 > 0.4) & (conf2 > 1.2)
    vaild = np.argwhere(logic).flatten()
    invaild = np.argwhere(~logic).flatten()
    ddmRet[invaild] = np.where(idx[invaild,1] < ddmDist//2, 0, ddmBank - 1)
    idx[:,1] += ddmRet * ddmDist
    org = np.array([cuberaw[idx[i,0], (ddmOrder*ddmDist+idx[i,1])%dBin,:].flatten() for i in range(len(idx))])
    return idx, snr, org, vaild

def RxBeamCfar(cuberaw, cube3d, cfarThld):
    rBin = cube3d.shape[0]
    dBin = cube3d.shape[1]
    aBin = cube3d.shape[2]
    mag = abs(cube3d)
    thld = np.sort(mag, axis=1)
    thld = thld[:,int(0.5*dBin),:]
    peak = np.zeros(mag.shape, dtype='b')
    for i in range(aBin):
        peak[:,:,i] = np.where(mag[:,:,i].T > thld[:,i]*(10**(cfarThld/20)), 1, 0).T
    peak = np.sum(peak, axis = -1)
    idx = np.argwhere(peak != 0)
    aziBins = np.argmax(mag[idx[:,0], idx[:,1], :],axis=-1)
    pwr = mag[idx[:,0], idx[:,1], aziBins]
    snr = 20*np.log10(pwr / thld[idx[:,0], aziBins])
    org = cuberaw[idx[i,0], idx[i,1], :]
    idx = np.c_[idx, aziBins]
    return idx, snr, org

def pySva(inArray: np.ndarray, axis=-1):
    """ Spatial Variant Apodization
        Instead of adding window in time domain, SVA is implemented in frequency domain.
        Suppose the original signal sequence is X(n), where n is in the range of 0 to N, and its FFT result is X(k), where k is in the range of 0 to K.
        Then after the FFT processing, SVA algorithm will be implemented in the following way
        a1 = 2 * real[X(k) * conj(X(k-1) + X(k+1))] / |X(k-1) + X(k+1)|^2
        Xw(k) = if a1 < 0:          X(k) 
        Xw(k) = if 0 <= a1 <= 1:    X(k) - a1 / 2 * (X(k-1) + X(k+1))
        Xw(k) = if a1 > 1:          X(k) - 1 / 2 * (X(k-1) + X(k+1))
    Args:
        inArray (np.ndarray): n dim fft result with no window
        axis (int or tuple): axis for sva. Defaults to -1, tuple for multi dim.

    Returns:
        np.ndarray: fft result after SVA
    """
    axis = axis if isinstance(axis, tuple) else (axis,)
    for i in axis:
        i += 0 if i >= 0 else inArray.ndim
        xk = inArray
        xkL = np.roll(inArray, -1, axis=i)
        xkR = np.roll(inArray,  1, axis=i)
        m = (xkL + xkR) / 2
        a_n = xk.real * m.real + xk.imag * m.imag
        a_d = m.real * m.real + m.imag * m.imag
        a = a_n / (1e-20+a_d)
        stage0 = (a_n < 0) | (a_d == 0)
        stage2 = a_n > a_d
        stage1 = ~(stage0 | stage2)
        xk = np.where(stage1, xk - a * m, xk)
        xk = np.where(stage2, xk - m, xk)
        inArray = xk
    return xk

def pyBLRC(steerVex:np.ndarray, sig:np.ndarray, rTH=1e-3, K=30):
    """
    BLRC based on Woodbury Matrix Identity
    reference: https://ieeexplore.ieee.org/document/10734934
    Parameters
    ----------
    steerVex : np.ndarray
        M x N, M is target beams, N is sensor number.
    sig : np.ndarray
        N vector.
    rTH : double, optional
        loop stop threshold. The default is 1e-3.
    K : double, optional
        max loops. The default is 30.

    Returns
    -------
    x : np.ndarray
        DESCRIPTION.
    """
    M, N = steerVex.shape
    y = sig
    A = steerVex.T
    AH = A.conj().T
    o2 = 0.1
    t = 1
    x = (AH @ y) / N
    for k in range(K):
        x2 = (abs(x))**2
        Q = 2 * o2 * t * np.diag(1/(1 + t * x2))
        Qinv = np.diag(1 / np.diag(Q))
        # ohm = o2 * np.linalg.inv(AH @ A + Q)
        ohm = o2 * (Qinv - Qinv @ AH @ np.linalg.inv(A @ Qinv @ AH+np.eye(N)) @ A @ Qinv)
        x = 1 / o2 * (ohm @ AH @ y)
        o2 = np.linalg.norm(y - A @ x, 2) / 2 / (N - np.trace(np.eye(M) - (Q @ ohm).real / o2))
        ohmm = np.diag(ohm).real
        x2 = (abs(x))**2
        k = t**2 * (ohmm + x2)**3 + 2 * t * x2 + ohmm + x2
        k /= (1 + t * (ohmm + x2)) ** 3
        _t = t
        t = M / 2 / np.sum(k)
        if abs(t - _t) < rTH:
            break
    return x

def ddmaCode2Order(code, banks=16):
    ddmOrder = []
    leakOrder = []
    for i in range(banks):
        ddmOrder.append(i) if (code>>i)&0x1==1 else leakOrder.append(i)
    pass
    return np.array(ddmOrder), np.array(leakOrder)

def fftshift(x, nfft, sftVal=None):
    if sftVal is None: sftVal = nfft // 2
    return np.where(x < sftVal, x, x - nfft)

def convolve2d_same(image: np.ndarray, kernel: np.ndarray, stride=1):
    """
    二维卷积操作，输出图像尺寸与输入图像相同

    参数:
        image (np.ndarray): 输入图像，形状为 (H, W)
        kernel (np.ndarray): 卷积核，形状为 (Kh, Kw)
        stride (int): 步长

    返回:
        output (np.ndarray): 卷积结果，形状为 (H, W)
    """
    H, W = image.shape
    Kh, Kw = kernel.shape

    # 计算需要的填充大小
    pad_h = (Kh - 1) // 2
    pad_w = (Kw - 1) // 2

    # 对输入图像进行填充
    image_padded = np.pad(image, ((pad_h, pad_h), (pad_w, pad_w)), mode='constant')

    # 使用滑动窗口提取所有窗口
    windows = np.lib.stride_tricks.sliding_window_view(image_padded, (Kh, Kw))[::stride, ::stride]

    # 计算卷积结果
    output = np.tensordot(windows, kernel, axes=((2, 3), (0, 1)))

    return output

def kalman_real(z, order=1, R=5.0, Q=0.01, ts=None):
    n = len(z)
    if ts is None:
        ts = np.arange(n, dtype=float)
    elif np.isscalar(ts):
        ts = np.full(n, ts, dtype=float)
    else:
        ts = np.asarray(ts, dtype=float)
    if order == 1:
        x = np.zeros((2, n))
        P = np.eye(2)
        H = np.array([[1, 0]])         # Observation matrix
        for i in range(n):
            if i > 0:
                dt = ts[i] - ts[i-1]
                F = np.array([[1, dt], [0, 1]])  # State transition with time interval
                x[:, i:i+1] = F @ x[:, i-1:i]
                Q_mat = Q * np.array([[dt**3/3, dt**2/2], [dt**2/2, dt]])
                P = F @ P @ F.T + Q_mat
            y = z[i] - (H @ x[:, i:i+1])[0, 0]
            S = H @ P @ H.T + R
            K = P @ H.T / S[0, 0]
            x[:, i:i+1] += K.reshape(-1, 1) * y
            P = (np.eye(2) - K.reshape(-1, 1) @ H) @ P
        return x[0, :], x[1, :], None
        
    elif order == 2:
        x = np.zeros((3, n))
        P = np.eye(3)
        H = np.array([[1, 0, 0]])
        for i in range(n):
            if i > 0:
                dt = ts[i] - ts[i-1]
                F = np.array([[1, dt, dt**2/2], [0, 1, dt], [0, 0, 1]])
                x[:, i:i+1] = F @ x[:, i-1:i]
                Q_mat = Q * np.array([[dt**5/20, dt**4/8, dt**3/6], 
                                    [dt**4/8, dt**3/3, dt**2/2], 
                                    [dt**3/6, dt**2/2, dt]])
                P = F @ P @ F.T + Q_mat
            y = z[i] - (H @ x[:, i:i+1])[0, 0]
            S = H @ P @ H.T + R
            K = P @ H.T / S[0, 0]
            x[:, i:i+1] += K.reshape(-1, 1) * y
            P = (np.eye(3) - K.reshape(-1, 1) @ H) @ P
        return x[0, :], x[1, :], x[2, :]
    else:
        raise ValueError("Order must be 1 or 2")

def kalman_cmplx(z, order=1, R=5.0, Q=0.01, ts=None):
    n = len(z)
    z_real = z.real
    z_imag = z.imag
    if ts is None:
        ts = np.arange(n, dtype=float)
    elif np.isscalar(ts):
        ts = np.full(n, ts, dtype=float)
    else:
        ts = np.asarray(ts, dtype=float)
    if order == 1:
        x = np.zeros((4, n))
        P = np.eye(4)
        H = np.array([[1, 0, 0, 0], [0, 0, 1, 0]])
        for i in range(n):
            if i > 0:
                dt = ts[i] - ts[i-1]
                F = np.array([[1, dt, 0, 0], [0, 1, 0, 0], [0, 0, 1, dt], [0, 0, 0, 1]])
                x[:, i:i+1] = F @ x[:, i-1:i]
                Q_block = Q * np.array([[dt**3/3, dt**2/2], [dt**2/2, dt]])
                Q_mat = np.block([[Q_block, np.zeros((2, 2))], [np.zeros((2, 2)), Q_block]])
                P = F @ P @ F.T + Q_mat
            y = np.array([z_real[i], z_imag[i]]) - (H @ x[:, i:i+1]).flatten()
            S = H @ P @ H.T + R * np.eye(2)
            K = P @ H.T @ np.linalg.inv(S)
            x[:, i:i+1] += K @ y.reshape(-1, 1)
            P = (np.eye(4) - K @ H) @ P
        filtered_complex = x[0, :] + 1j * x[2, :]
        vel = x[1, :] + 1j * x[3, :]
        return filtered_complex, vel, None
    elif order == 2:
        x = np.zeros((6, n))
        P = np.eye(6)
        H = np.array([[1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0]])
        for i in range(n):
            if i > 0:
                dt = ts[i] - ts[i-1]
                F = np.array([[1, dt, dt**2/2, 0, 0, 0], [0, 1, dt, 0, 0, 0], [0, 0, 1, 0, 0, 0],
                              [0, 0, 0, 1, dt, dt**2/2], [0, 0, 0, 0, 1, dt], [0, 0, 0, 0, 0, 1]])
                x[:, i:i+1] = F @ x[:, i-1:i]
                Q_block = Q * np.array([[dt**5/20, dt**4/8, dt**3/6], 
                                        [dt**4/8, dt**3/3, dt**2/2], 
                                        [dt**3/6, dt**2/2, dt]])
                Q_mat = np.block([[Q_block, np.zeros((3, 3))], [np.zeros((3, 3)), Q_block]])
                P = F @ P @ F.T + Q_mat
            y = np.array([z_real[i], z_imag[i]]) - (H @ x[:, i:i+1]).flatten()
            S = H @ P @ H.T + R * np.eye(2)
            K = P @ H.T @ np.linalg.inv(S)
            x[:, i:i+1] += K @ y.reshape(-1, 1)
            P = (np.eye(6) - K @ H) @ P
        filtered_complex = x[0, :] + 1j * x[3, :]
        vel = x[1, :] + 1j * x[4, :]
        accel = x[2, :] + 1j * x[5, :]
        return filtered_complex, vel, accel
    else:
        raise ValueError("Order must be 1 or 2")

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import time
    def snr(x):
        noise = 20*np.log10(np.mean(x[:, 100:-100]))
        sig = 20*np.log10(np.max(x))
        print(f"sig:{sig}, noise:{noise}, snr:{sig-noise}")
    RBIN = 2048
    VBIN = 2048
    BENCH_TIME = 10
    WINTYPE = np.hamming
    
    sig = (np.random.random([RBIN, VBIN]) - 0.5) + 1j * (np.random.random([RBIN, VBIN]) - 0.5)
    n = 8
    k = 1/n
    for i in range(2*n):
        sig += 1 * np.outer(np.exp((20*(1+i) + k*((i%n-n//2)))*1j*np.arange(RBIN)/RBIN*2*np.pi),
                            np.exp((20*(1+i) + k*((i%n-n//2)))*1j*np.arange(VBIN)/VBIN*2*np.pi))
        break
    win = np.outer(WINTYPE(RBIN), WINTYPE(VBIN))
    fft2d = np.fft.fft2(sig)
    fft2d = fft2d.astype(np.complex64)
    start = time.time()
    for i in range(BENCH_TIME):
        pySva(fft2d, axis=0)
    print(f'range sva cost {(time.time()-start)/BENCH_TIME} s')
    start = time.time()
    for i in range(BENCH_TIME):
        pySva(fft2d, axis=1)
    print(f'doppler sva cost {(time.time()-start)/BENCH_TIME} s')
    start = time.time()
    for i in range(BENCH_TIME):
        pySva(fft2d, axis=(0, 1))
    print(f'2d sva cost {(time.time()-start)/BENCH_TIME} s')
    # plt.figure()
    # plt.title('2Dfft Test')
    raw = abs(np.fft.fft2(sig))
    rawWin = abs(np.fft.fft2(sig*win))
    rawSva = abs(pySva(np.fft.fft2(sig), axis=(0, 1)))
    plt.plot(20*np.log10(raw[:,[0,300]]), label='rect win')
    plt.plot(20*np.log10(rawWin[:,[0,300]]), label='hamming win')
    plt.plot(20*np.log10(rawSva[:,[0,300]]), label='SVA')
    snr(raw);snr(rawWin);snr(rawSva)
    # plt.legend()
    # plt.show()
    #%% DFT test
    N = 194
    nfft = 192
    p = np.r_[0,np.arange(2,192),193]
    sig = 15*np.exp(50.25j*np.arange(N)/N*2*np.pi) + (np.random.random(N) - 0.5)
    sig = sig[p]
    az = np.arange(nfft)-nfft//2
    steerVec = np.outer(az/nfft*2*np.pi, p)
    steerVec = np.exp(1j*steerVec)
    rawfft = np.dot(steerVec, sig)
    rawfft = np.round(rawfft)
    win = WINTYPE(194)
    win = np.r_[win[0], win[2:-2], win[-1]]
    plt.figure()
    plt.title('DFT Test')
    plt.plot(20*np.log10(abs(np.dot(steerVec, sig))), label='rect win')
    plt.plot(20*np.log10(abs(np.dot(steerVec, sig*win))), label='hamming win')
    plt.plot(20*np.log10(abs(pySva(np.dot(steerVec, sig)))), label='SVA')
    plt.plot(20*np.log10(abs(pySva(rawfft))), label='fft')
    plt.legend()
    plt.show()
    
#%%
    N = 32
    sig = 15*np.exp(50.25j*np.arange(N)/N*2*np.pi) + (np.random.random(N) - 0.5)
    x = sig.copy()
    x[[1,2,4,5,8,10,12,13,14]] = 0
    plt.figure()
    plt.title('DFT Test')
    plt.plot(20*np.log10(abs(np.fft.fft(sig))), label='full fft')
    plt.plot(20*np.log10(abs(np.fft.fft(x))), label='zero-padding fft')
    